// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/vector_function.h>
#include <pollux/functions/prestosql/types/bing_tile_type.h>
#include <pollux/functions/prestosql/types/geometry_type.h>
#include <pollux/functions/prestosql/types/hyper_log_log_type.h>
#include <pollux/functions/prestosql/types/ip_address_type.h>
#include <pollux/functions/prestosql/types/ip_prefix_type.h>
#include <pollux/functions/prestosql/types/json_type.h>
#include <pollux/functions/prestosql/types/tdigest_type.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>
#include <pollux/functions/prestosql/types/uuid_type.h>

namespace kumo::pollux::functions {
    namespace {
        std::string typeName(const TypePtr &type) {
            switch (type->kind()) {
                case TypeKind::BOOLEAN:
                    return "boolean";
                case TypeKind::TINYINT:
                    return "tinyint";
                case TypeKind::SMALLINT:
                    return "smallint";
                case TypeKind::INTEGER:
                    if (type->isDate()) {
                        return "date";
                    }
                    if (type->isIntervalYearMonth()) {
                        return "interval year to month";
                    }
                    return "integer";
                case TypeKind::BIGINT:
                    if (isTimestampWithTimeZoneType(type)) {
                        return "timestamp with time zone";
                    }
                    if (type->isIntervalDayTime()) {
                        return "interval day to second";
                    }
                    if (type->isDecimal()) {
                        const auto &shortDecimal = type->asShortDecimal();
                        return fmt::format(
                            "decimal({},{})", shortDecimal.precision(), shortDecimal.scale());
                    }
                    if (isBingTileType(type)) {
                        return "bingtile";
                    }
                    return "bigint";
                case TypeKind::HUGEINT: {
                    if (isUuidType(type)) {
                        return "uuid";
                    } else if (isIPAddressType(type)) {
                        return "ipaddress";
                    }
                    POLLUX_USER_CHECK(
                        type->isDecimal(),
                        "Expected decimal type. Got: {}",
                        type->toString());
                    const auto &longDecimal = type->asLongDecimal();
                    return fmt::format(
                        "decimal({},{})", longDecimal.precision(), longDecimal.scale());
                }
                case TypeKind::REAL:
                    return "real";
                case TypeKind::DOUBLE:
                    return "double";
                case TypeKind::VARCHAR:
                    if (isJsonType(type)) {
                        return "json";
                    }
                    return "varchar";
                case TypeKind::VARBINARY:
                    if (isHyperLogLogType(type)) {
                        return "HyperLogLog";
                    }
                    if (isGeometryType(type)) {
                        return "geometry";
                    }
                    if (*type == *TDIGEST(DOUBLE())) {
                        return "tdigest(double)";
                    }
                    return "varbinary";
                case TypeKind::TIMESTAMP:
                    return "timestamp";
                case TypeKind::ARRAY:
                    return fmt::format("array({})", typeName(type->childAt(0)));
                case TypeKind::MAP:
                    return fmt::format(
                        "map({}, {})",
                        typeName(type->childAt(0)),
                        typeName(type->childAt(1)));
                case TypeKind::ROW: {
                    if (isIPPrefixType(type)) {
                        return "ipprefix";
                    }
                    const auto &rowType = type->as_row();
                    std::ostringstream out;
                    out << "row(";
                    for (auto i = 0; i < type->size(); ++i) {
                        if (i > 0) {
                            out << ", ";
                        }
                        if (!rowType.nameOf(i).empty()) {
                            out << "\"" << rowType.nameOf(i) << "\" ";
                        }
                        out << typeName(type->childAt(i));
                    }
                    out << ")";
                    return out.str();
                }
                case TypeKind::UNKNOWN:
                    return "unknown";
                default:
                    POLLUX_UNSUPPORTED("Unsupported type: {}", type->toString());
            }
        }

        class TypeOfFunction : public exec::VectorFunction {
        public:
            TypeOfFunction(const TypePtr &type) : typeName_{typeName(type)} {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                auto localResult = BaseVector::create_constant(
                    VARCHAR(), typeName_, rows.size(), context.pool());
                context.moveOrCopyResult(localResult, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // T -> varchar
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("varchar")
                    .argumentType("T")
                    .build()
                };
            }

            static std::shared_ptr<exec::VectorFunction> create(
                const std::string & /*name*/,
                const std::vector<exec::VectorFunctionArg> &inputArgs,
                const core::QueryConfig & /*config*/) {
                try {
                    return std::make_shared<TypeOfFunction>(inputArgs[0].type);
                } catch (...) {
                    return std::make_shared<exec::AlwaysFailingVectorFunction>(
                        std::current_exception());
                }
            }

        private:
            const std::string typeName_;
        };
    } // namespace

    POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION_WITH_METADATA(
        udf_typeof,
        TypeOfFunction::signatures(),
        exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
        TypeOfFunction::create);
} // namespace kumo::pollux::functions
